#-------------------------------------------------------------------------------
# Name:        MAPP Analysis
# Purpose:     Provides functions for protein substitutions analysis
#              via MAPP method.
#              Don't forget to set program paths in settings openedInFile.
#
# Author:      Bc. Onrej Pelikan
#
# Created:     25.03.2013
#-------------------------------------------------------------------------------

import re
import settings
import logging
from xml.etree import ElementTree
from blasttofasta import convertXMLtoFASTA
from subprocess import Popen, PIPE, STDOUT
import os
import uuid
import operator

#FILES USED BY PROGRAMS
fileid = uuid.uuid4()

msafile = settings.msafile + str(fileid)
treefile = settings.treefile + str(fileid)
mappfile = settings.mappfile + str(fileid)


#MAPP HEADER
mappHeader = "Mutation\tPosition\tColumn Score\tColumn p-value\tAlignment\tGap Weight\tOver Gap Weight Threshold\tHydropathy\tPolarity\tCharge\tVolume\tFree energy alpha\tFree energy beta\tA\tC\tD\tE\tF\tG\tH\tI\tK\tL\tM\tN\tP\tQ\tR\tS\tT\tV\tW\tY\tA\tC\tD\tE\tF\tG\tH\tI\tK\tL\tM\tN\tP\tQ\tR\tS\tT\tV\tW\tY\tGood Amino Acids\tBad Amino Acids"
mappHeaderList = mappHeader.split("\t")

#P-VALUE THRESHOLD
pvalueThreshold = 0.1 #p-value less then threshold means deleterious
NEUTRAL = "Neutral"
DELETERIOUS = "Deleterious"
NA = "N/A"

# MAPP OUTPUT FILE OFFSETS
aalist = ["A","C","D","E","F","G","H","I","K","L","M","N","P","Q","R","S","T","V","W","Y"]
nameOffset = 0
naOffset = 6
scoreOffset = 13
pvalueOffset = 33
alignmentOffset = 4
hydroOffset = 7
polarityOffset = 8
chargeOffset = 9
volumeOffset = 10
alfaEnergyOffset = 11
betaEnergyOffset = 12



logger = logging.getLogger(__name__)


# EXCEPTION IN ANLYSIS PROCESS
class MappException(Exception):
    """ This exception contains .header of error and error .description."""
    def __init__(self, errorHeader, errorDescription):
        self.header = errorHeader
        self.description = errorDescription

#CLASS OF PROGRAM VARIABLES - Its runPrograms's method input.
class ProgramVariables:
    """ "sequence" - seqence of amino acid letters
        "mutations" - list of tuples [(letter, position, new letter), ... ]
        "database_type" - type of clustal database "nr", "refseq", "swissprot"
        "error" - error for blast e.g. 1E-12 or 0.0001
        "treeMethod" - tree construction method "nj", "ml" (neighbour joining, maximum likelihood)
        "num_alignments" - maximum number of given sequeces
        "inclusion_ethresh" - threshold for next psi-blast iteration e.g. 0.002
        "num_iterations" - number of psi blast iteration
    """
    semphyParam = {"nj" : ["--jtt", "-J", "-H"],
                    "ml" : ["--jtt", "-S", "--posteriorDTME"]}

    def __init__(self, sequence, mutations, database_type = "nr", blastError = 0.0001, treeMethod = "nj",
                    num_alignments = 399,
                    inclusion_ethresh = 0.002,
                    num_iterations = 3):
        self.sequence = sequence
        self.num_alignments = str(num_alignments)
        self.mutations = mutations
        self.database_type = database_type
        self.blastError = str(blastError)
        self.inclusion_ethresh=str(inclusion_ethresh)
        self.num_iterations=str(num_iterations)
        self.treeMethod = treeMethod

    def __str__(self):
        mutations = ", ".join(("".join(mutation) for mutation in self.mutations))
        return "sequence={0}, mutations={1}, database_type={2}, blastError={3}, treeMethod={4}".format(
                    self.sequence[0:6] + "..." if (len(self.sequence) > 7) else self.sequence,
                    mutations[0:6] + "..." if (len(mutations) > 7) else mutations,
                    self.database_type,
                    str(self.blastError),
                    self.treeMethod)



# MAJOR FUNCTION FOR ANALYSIS
def runPrograms(progVariables, mapponly=False):
    """ Runs programs needed for anlysis.
        Parameter progVariables is instance of ProgramVariables class.
        Parameter mapponly skip blast, clustal and semphy. Files that created clustal and semphy
            must be created manualy before.
        Returns exactly what returns processMappOutputMethod.
        Raises MappException when error occures. """
    try:
        logger.info("Program variables for varibench test: {0}".format(str(progVariables)))
        if not mapponly:
            blastout = runBlast(progVariables.sequence, progVariables.database_type, progVariables.blastError,
                                 progVariables.num_alignments, progVariables.inclusion_ethresh, progVariables.num_iterations)
            fasta = convertBlast(blastout)
            logger.info("Blast found {0} sequences.".format(fasta.count(">")))
            #convert fasta output to mapp friendly input according to mapp documentation
            fasta = re.sub("[^\w\n\r>_]", "_", fasta).replace("\r", "\n").replace("\n\n", "\n").replace("_\n","\n")
            runOmegaAndWriteToFile(msafile, fasta)
        mappout = runMapp(treefile, msafile).replace(",", ".")
        newMutations = correctMutationPositions(progVariables.mutations, msafile)
        openedInFile = open(mappfile, "w")
        openedInFile.write(mappout)
        openedInFile.close()
    except MappException:
        raise

    removeFiles()

    return processMappOutput(progVariables.mutations, newMutations, mappout)


def processMappOutput(oldMutations, newMutations, mappout):
    """ Takes mapp output and pick out only line that contatins desirable mutations.
        Returns list of mapp lines where values are delimited by tabs. Fisrst value on the line
            is mutation name and others are orignal mapp values.
        Parameter newMutations is list of tuples [(letter, position, new letter), ... ]
        Parameter oldMutations gives old names of mutations before correction."""

    lines = mappout.splitlines()
    processedMutations = list()

    for i in range(len(newMutations)):
        if newMutations[i][1] is None:
            continue
        mappline = lines[int(newMutations[i][1])] # mutations starts from 1, and mapplines starts from 1 too (line 0 is header)
        mutationName = "".join(str(m) for m in oldMutations[i])
        processedMutations.append(mutationName + "\t" + mappline)

    return processedMutations


#RUN PROGRAMS
def runBlast(sequence, database, error, num_alignments, inclusionEthresh, numIterations):
    """ sequence is string of amino acid letters
        database parameter contains a used database (see "ProgramVariables.database_type").
        error parameter contains number in format 0.0001 or 1E-4.
        num_alignments parameter contains number of returned sequences.
        inclusionEthresh parameter contains threshold for next psi-blast iteration
        numIterations parameter contains number of psi-blast iteration
        http://www.ncbi.nlm.nih.gov/books/NBK1763/ """
    logger.info("Blast is starting.")
    logger.info(getBlastRunCmd(database, error, num_alignments, inclusionEthresh, numIterations))
    program = Popen(getBlastRunCmd(database, error, num_alignments, inclusionEthresh, numIterations), stdout=PIPE, stdin=PIPE, stderr=PIPE)
    blastout = program.communicate(input=">\n" + sequence)
    logger.info("Blast has ended.")
    if program.returncode != 0:
        logger.warning("Psiblast program return with error: " + blastout[1])
        raise MappException("Psiblast program return with error.", blastout[1])
    return blastout[0]

def convertBlast(blast):
    fasta = convertXMLtoFASTA(blast)
    return fasta

def runOmegaAndWriteToFile(msafile, fasta):
    """ Takes fasta and clustal oomega output is written to msafile
    http://www.clustal.org/omega/README """
    logger.info("Clustal omega is starting.")
    logger.info(getClustalRunCmd(msafile))
    program = Popen(getClustalRunCmd(msafile), stdin=PIPE)
    program.communicate(input=fasta)
    logger.info("Clustal omega has ended.")
    if program.returncode != 0:
        logger.warning("Clustal Omega program ended with error!")
        raise MappException("Clustal Omega program ended with error!", "Exit code:" + str(program.returncode))

def runMapp(treefile, msafile):
    """ Tales treefile and msafile, provides mapp analysis and
    returns raw mapp program output """
    logger.info("Mapp is starting.")
    logger.info(getMappRunCmd(treefile, msafile))
    program = Popen(getMappRunCmd(treefile, msafile), stdin=PIPE, stdout=PIPE, stderr=PIPE)
    mappout = program.communicate()
    logger.info("Mapp has ended.")
    if program.returncode != 0 or len(mappout[1]) != 0:
        logger.warning("Mapp tool exited with error! " + mappout[1])
        raise MappException("Mapp tool exited with error!", "Exit code: " + str(program.returncode) + "\n" + mappout[1])

    return mappout[0]


# PEOPEN ARGS
def getBlastRunCmd(database, evalue, num_alignments, inclusionEthresh, numIterations):
    #return [settings.blast, "-remote", "-db", database, "-outfmt", "5", "-evalue", evalue, "-num_alignments", str(num_alignments), "-inclusion_ethresh", inclusionEthresh, "-num_iterations", numIterations]
    command = [settings.blast, "-outfmt", "5", "-evalue", evalue, "-num_alignments", str(num_alignments), "-inclusion_ethresh", inclusionEthresh]
    if settings.remote:
        command.extend(["-remote", "-db", database])
    else:
        command.extend(["-db", os.path.join(settings.databaseDir, database), "-num_iterations", numIterations])
    return command

def getClustalRunCmd(msafile):
    return [settings.clustalomega, "--infile=-", "-o", msafile, "--outfmt=fa", "--auto", "--force"]

def getMappRunCmd(treefile, msafile):
    return ["java", "-jar", settings.mapp, "-f", msafile, "-t", treefile]


def parseMutations(mutations):
    """ This function returns list of tuples [(letter, position, new letter), ... ]
            Position is string (not int).
         It searches letterNumbersLetter regular expression in mutations string."""
    return re.compile("([a-zA-Z])(\d+)([a-zA-Z])").findall(mutations)



def checkMutations(sequence, mutations):
    """ Check if first letter of mutation matches mutation position in sequence. """
    error = False
    for mutation in mutations:
        if sequence[int(mutation[1])-1] != mutation[0]:
            print("Mutation {0} doesn't match. In sequence occurs: {1}".format(str(mutation), sequence[int(mutation[1])-1]))
            error = True
    if error:
        return False
    else:
        return True


def correctMutationPositions(oldMutations, msafile):
    """ Original list of mutations not suppose spaces in alignemnt. Thus, positions
            have to be corrected.
        Parameter oldMutations is list of tupples [(letter, position, letter), ... ]
        Parameter msafile is path to the multiple sequence alignment openedInFile.
        Returns list of tupples containing new positions according to the msa. Wrpmg
            mutations are replaced by (None, None, None).
        """
    openedInFile = open(msafile, "r")
    sequence = re.compile(">.+\n(([a-zA-Z\-]+\n)+)").match(openedInFile.read()).group(1).replace("\n","")
    openedInFile.close()
    newMutations = []
    for mutation in oldMutations:
        oldPosition = int(mutation[1])
        letterCount = 0
        newPosition = 0
        while letterCount < oldPosition:
            newPosition += 1
            if re.match("[a-zA-Z]", sequence[newPosition-1]):
                letterCount += 1
        if (sequence[newPosition-1] != mutation[0]):
            newMutation = (None, None, None)
            print("Seq: " + sequence[:15] + "..., mut: " + str(mutation) + ", ERROR. Mutation letter is wrong. Mutation is excluded." )
        else:
            newMutation = (mutation[0], newPosition, mutation[2])

        newMutations.append(newMutation)


    return newMutations


def transformAlignmentString(alignmentString):
    """ Transforms alignment string from mapp output to readable form.
            From "AAAAAACCDE" makes "6xA 2xC 1xD 1xE".
        Returns new alignment string.
    """
    aadict = dict(zip(aalist, [0]*len(aalist)))
    aadict["-"] = 0
    aadict["X"] = 0
    alignmentString = alignmentString.lstrip("'")
    for letter in alignmentString:
        aadict[letter] += 1
    sortedAlignment = sorted(aadict.iteritems(), key=operator.itemgetter(1), reverse=True)
    align = ("{c}x{l}".format(c=count, l=letter) for letter, count in sortedAlignment if count > 0)
    return " ".join(align)


def evalMutations(lines):
    """ This method reads and process mapp output lines.
        Parameter lines is iterable item contains lines from runProgram methods.
            (It can be even opened openedInFile object)
        Returns dictionary {"neutral" - netural mutation count
                            "delterious" - deleterious mutation count
                            "na" - N/A mutation count
                            "mutations" - list of [{"name"
                                                   "effect" - module variable DELETERIOIUS, NEUTRAL, NA
                                                   "p_value"
                                                   "prediction_score"
                                                   "alignment"
                                                    }, ...]
                            }
    """
    neutral = 0
    deleterious = 0
    noteval = 0
    evaluatedLst = []
    for line in lines:
        if (line == "\n"):
            break
        line = line.split("\t")

        newLetter = parseMutations(line[nameOffset])[0][2] #parseMutatons returns [(letter, position, new letter), ... ]
        alignment = transformAlignmentString(line[alignmentOffset])

        if (line[naOffset] == "N"): # values are defined
            p_value = float(line[pvalueOffset + aalist.index(newLetter)])
            pred_score = float(line[scoreOffset+aalist.index(newLetter)])
            if (p_value > pvalueThreshold):
                effect = NEUTRAL
                neutral += 1
            else:
                effect = DELETERIOUS
                deleterious += 1
        else:
            effect = NA
            p_value = float("nan")
            pred_score = float("nan")
            noteval += 1
        evaluatedLst.append({"name" : line[nameOffset],
                                "effect": effect,
                                "p_value": p_value,
                                "prediction_score": pred_score,
                                "alignment": alignment})

    return {"neutral": neutral,
            "deleterious": deleterious,
            "na" : noteval,
            "mutations": evaluatedLst}


def removeFiles():
    if settings.removefiles:
        os.remove(msafile)
        os.remove(treefile)
        os.remove(mappfile)


# TODO predelat tak, aby jeden radek rozedelil do slovniku (nazvy podle mappu) a
#       a tuhle metodu dat do utils modulu.
##def processMappOutput(mutations, mappout):
##    """ Takes Mapp output and process it.
##        Returns list of tuples:
##            [(mutation name, effect, aa pvalue, aa score, alignment, hydropathy, polarity, charge, volume, free energy alpha, free energy beta), ...]
##            of types (string, float, float    , float   , string   , float     , float   , float , float , float            , float           ) """
##
##    mappout = [line.split("\t") for line in mappout.splitlines()]
##
##    processedMutations = list()
##    for mutation in mutations:
##        offset = aalist.index(mutation[2])
##        mappline = mappout[int(mutation[1])]
##
##        mutationName = "".join(mutation)
##        score = float(mappline[scoreOffset + offset])
##        pvalue = float(mappline[pvalueOffset + offset])
##        effect = "deleterious" if pvalue < pvalueThreshold else "neutral"
##        alignment = mappline[alignmentOffset]
##        alignment = parseAlignment(alignment)
##        processedMutations.append((mutationName, effect, pvalue, score, alignment,
##                                    float(mappline[polarityOffset]), float(mappline[chargeOffset]),
##                                    float(mappline[volumeOffset]), float(mappline[alfaEnergyOffset]),
##                                    float(mappline[betaEnergyOffset])))
##
##    return processedMutations




